!> multigroup operators and solution+rhs vectors
!!
!!  This module contains the routine for constructing the
!!  operator matrices and corresponding vectors needed to
!!  solve the linear system.
!!
!!  J. Roberts, 8/1/2010
module matx_const

    use class_input, only: indat
    use class_dat, only: rhsdat,coefdat,flux

    implicit none

    contains

!************************************************************
!> given user input, construct the K matrix and other vectors
!!  and finishes filling out the input class
!************************************************************
    subroutine setk(A,inp,rhs,dat)

        implicit none
#include "finclude/petsc.h90"

        Mat, dimension(:), pointer, intent(inout)  :: A
        type(indat), intent(inout)                 :: inp
        type(rhsdat), intent(inout)                :: rhs
        type(coefdat), intent(inout)               :: dat
        ! local variables
        PetscInt :: i, j, ind_str, ind_end, k, m, g, gg, nx, ny,       &
                    m_L, m_R, m_C, m_B, m_T, lenk, beta, beta1, beta2, &
                    rownum, colnum, rowidx(1), colidx(5), numg, ierr
        PetscScalar, dimension(:), pointer       :: dx, dy, dv
        PetscScalar                              :: A_C, A_L, A_R, A_B, A_T,poo
        PetscScalar, dimension(:,:), allocatable :: xsd, alphaL, alphaR, alphaB, alphaT
        PetscInt, dimension(:), allocatable      :: cix, ciy
        ! temporary pointers to input
        PetscScalar, dimension(:,:), pointer     :: dc=>null(), sr=>null(), ab=>null(), &
                                                    ns=>null(), xi=>null(), sa=>null()
        PetscInt, dimension(:,:), pointer        :: mt=>null()
        PetscScalar, dimension(:,:,:), pointer   :: sc=>null(), src=>null()

        ! point to the input for cleaner code
        numg = inp%numg
        dc        => inp%dc
        sr        => inp%sr
        ab        => inp%ab
        ns        => inp%ns
        xi        => inp%xi
        sc        => inp%sc
        sa        => inp%sa
        src       => inp%src
        mt        => inp%mt

        ! compute dx and dy vectors
        allocate(dx(inp%nxfm), dy(inp%nyfm), dv(inp%nyfm*inp%nxfm))
        ! coarse mesh index for fine mesh
        allocate(cix(inp%nxfm),ciy(inp%nyfm))

        j = 0
        do i = 1, inp%nxcm
            ind_str = j+1
            ind_end = j+inp%xfm(i)
            do gg = ind_str, ind_end
                dx(gg)  = (inp%xcm(i+1)-inp%xcm(i))/inp%xfm(i)
                cix(gg) = i
            end do
            j = sum(inp%xfm(1:i))
        end do
        j = 0
        do i = 1, inp%nycm
            ind_str = j+1
            ind_end = j+inp%xfm(i)
            do gg = ind_str, ind_end
                dy(gg)  = (inp%ycm(i+1)-inp%ycm(i))/inp%yfm(i)
                ciy(gg) = i
            end do
            j = sum(inp%yfm(1:i))
        end do
        lenk = inp%nxfm*inp%nyfm ! number of rows per group

        ! allocate the rhs and dat vectors
        allocate( rhs%fixedSource(numg), rhs%nuSigF(numg), rhs%xiSpect(numg), rhs%sigScatter(numg,numg) )
        allocate( dat%sigAbs(numg) )
        call VecCreate( PETSC_COMM_WORLD, rhs%fixedSource(1), ierr )
        call VecSetSizes( rhs%fixedSource(1), PETSC_DECIDE, lenk, ierr )
        call VecSetFromOptions( rhs%fixedSource(1), ierr )
        call VecDuplicate( rhs%fixedSource(1), rhs%nuSigF(1) , ierr )
        call VecDuplicate( rhs%fixedSource(1), rhs%xiSpect(1), ierr )
        call VecDuplicate( rhs%fixedSource(1), dat%sigAbs(1) , ierr )
        ! allocate other vectors
        do g = 1, numg
            do gg = 1, numg
                call VecDuplicate( rhs%fixedSource(1), rhs%sigScatter(g,gg), ierr)
            end do
        end do
        do g = 2, numg
            call VecDuplicate( rhs%fixedSource(1),rhs%fixedSource(g),ierr)
            call VecDuplicate( rhs%fixedSource(1),rhs%nuSigF(g)     ,ierr)
            call VecDuplicate( rhs%fixedSource(1),rhs%xiSpect(g)    ,ierr)
            call VecDuplicate( rhs%fixedSource(1),dat%sigAbs(g)     ,ierr)
        end do

        call VecCreate( PETSC_COMM_WORLD, inp%dx, ierr )
        call VecSetSizes( inp%dx, PETSC_DECIDE, inp%nxfm, ierr )
        call VecSetFromOptions( inp%dx, ierr )
        call VecCreate( PETSC_COMM_WORLD, inp%dy, ierr )
        call VecSetSizes( inp%dy, PETSC_DECIDE, inp%nxfm, ierr )
        call VecSetFromOptions( inp%dy, ierr )
        call VecDuplicate( rhs%fixedSource(1), inp%dv , ierr )

        do i = 1, inp%nxfm
            call VecSetValue( inp%dx, i-1, dx(i), INSERT_VALUES, ierr )
            do j = 1, inp%nyfm
                if (i.eq.1) call VecSetValue( inp%dy, j-1, dy(j), INSERT_VALUES, ierr )
                k     = i+(j-1)*( inp%nxfm )
                dv(k) = dx(i) * dy(j)
                call VecSetValue( inp%dv, k-1, dv(k), INSERT_VALUES, ierr )
            end do
        end do
        call VecAssemblyBegin(inp%dx,ierr)
        call VecAssemblyEnd(inp%dx,ierr)
        call VecAssemblyBegin(inp%dy,ierr)
        call VecAssemblyEnd(inp%dy,ierr)
        call VecAssemblyBegin(inp%dv,ierr)
        call VecAssemblyEnd(inp%dv,ierr)

        ! need to specify alpha better either a constant, or by coarse mesh, or
        ! by a sum of DLPs ( c0*DLP(0) + c1*DLP(1) + ... cN*DLP(N) )

        if (inp%bcl .eq. 2) then
            allocate( alphaL(inp%nyfm,numg) )
            alphaL = 1.0
            !alphaL = 2*bcl!rand(inp%nyfm,numg)!'!ones(inp%nyfm,numg)
        end if
        if (inp%bcr .eq. 2) then
            allocate( alphaR(inp%nyfm,numg) )
            alphaR = 1.0
            !alphaR = 2*IBSR!ones(inp%nyfm,numg)
        end if
        if (inp%bcb .eq. 2) then
            allocate( alphaB(inp%nxfm,numg) )
            alphaB = 1.0
            !alphaB = 2*IBSB!ones(nxfm,numg)
        end if
        if (inp%bct .eq. 2) then
            allocate( alphaR(inp%nxfm,numg) )
            alphaT = 1.0
            !alphaT = 2*IBST!ones(nxfm,numg)
        end if

        do g = 1, numg ! group loop

            ! interior coefficients
            do i = 2,inp%nxfm-1
                do j = 2,inp%nyfm-1
                    ! materials
                    m_L = mt( cix(i-1), ciy(j  ) )
                    m_R = mt( cix(i+1), ciy(j  ) )
                    m_B = mt( cix(i  ), ciy(j-1) )
                    m_T = mt( cix(i  ), ciy(j+1) )
                    m_C = mt( cix(i  ), ciy(j  ) )
                    ! matrix row index
                    k = i+(j-1)*( inp%nxfm )
                    ! define matrix values
                    A_L = -2.0*dc(m_L,g)*dc(m_C,g)*dy(j)/(dx(i-1)*dc(m_C,g)+dx(i)*dc(m_L,g))
                    A_R = -2.0*dc(m_R,g)*dc(m_C,g)*dy(j)/(dx(i+1)*dc(m_C,g)+dx(i)*dc(m_R,g))
                    A_B = -2.0*dc(m_B,g)*dc(m_C,g)*dx(i)/(dy(j-1)*dc(m_C,g)+dy(j)*dc(m_B,g))
                    A_T = -2.0*dc(m_T,g)*dc(m_C,g)*dx(i)/(dy(j+1)*dc(m_C,g)+dy(j)*dc(m_T,g))
                    A_C = dv(k)*sr(m_C,g)-(A_L+A_R+A_B+A_T)

                    call SetValues( A, inp, rhs, dat, k, lenk, g, (/A_L,A_R,A_B,A_T,A_C/), &
                                    (/1,1,1,1/), dv(k), cix(i), ciy(j), m_C )

                end do
            end do

            ! left and right edges
            do j = 2, inp%nyfm-1

                !----------------------------------------------
                ! LEFT EDGE
                !----------------------------------------------

                i = 1
                k = i+(j-1)*( inp%nxfm )

                if (inp%bcl .eq. 1) then
                    beta=1
                else
                    beta=0
                end if

                ! materials
                m_R = mt( cix(i+1), ciy(j  ) )
                m_B = mt( cix(i  ), ciy(j-1) )
                m_T = mt( cix(i  ), ciy(j+1) )
                m_C = mt( cix(i  ), ciy(j  ) )

                A_R = -2.0*dc(m_R,g)*dc(m_C,g)*dy(j)/(dx(i+1)*dc(m_C,g)+dx(i)*dc(m_R,g))
                A_B = -2.0*dc(m_B,g)*dc(m_C,g)*dx(i)/(dy(j-1)*dc(m_C,g)+dy(j)*dc(m_B,g))
                A_T = -2.0*dc(m_T,g)*dc(m_C,g)*dx(i)/(dy(j+1)*dc(m_C,g)+dy(j)*dc(m_T,g))
                A_C = dv(k)*sr(m_C,g) - (A_R+A_B+A_T) &
                      + dc(m_C,g)*dy(j)*2*(1-beta)/(4*(1+beta)*dc(m_C,g)+dx(i)*(1-beta))

                call SetValues( A, inp, rhs, dat, k, lenk, g, (/A_L,A_R,A_B,A_T,A_C/), &
                                (/0,1,1,1/), dv(k), cix(i), ciy(j), m_C )

                !----------------------------------------------
                ! RIGHT EDGE
                !----------------------------------------------

                i = inp%nxfm
                k = i+(j-1)*( inp%nxfm )

                if (inp%bcr.eq.1) then
                    beta=1
                else
                    beta=0
                end if

                ! materials
                m_L = mt( cix(i-1), ciy(j  ) )
                m_B = mt( cix(i  ), ciy(j-1) )
                m_T = mt( cix(i  ), ciy(j+1) )
                m_C = mt( cix(i  ), ciy(j  ) )

                A_L    = -2.0*dc(m_L,g)*dc(m_C,g)*dy(j)/(dx(i-1)*dc(m_C,g)+dx(i)*dc(m_L,g))
                A_B    = -2.0*dc(m_B,g)*dc(m_C,g)*dx(i)/(dy(j-1)*dc(m_C,g)+dy(j)*dc(m_B,g))
                A_T    = -2.0*dc(m_T,g)*dc(m_C,g)*dx(i)/(dy(j+1)*dc(m_C,g)+dy(j)*dc(m_T,g))
                A_C    = dv(k)*sr(m_C,g) - (A_L+A_B+A_T) &
                         + dc(m_C,g)*dy(j)*2*(1-beta)/(4*(1+beta)*dc(m_C,g)+dx(i)*(1-beta))

                call SetValues( A, inp, rhs, dat, k, lenk, g, (/A_L,A_R,A_B,A_T,A_C/), &
                                (/1,0,1,1/), dv(k), cix(i), ciy(j), m_C )

            end do

            ! top and bottom edges
            do i = 2, inp%nxfm-1

                !----------------------------------------------
                ! BOTTOM EDGE
                !----------------------------------------------

                j = 1
                k = i+(j-1)*( inp%nxfm )

                if (inp%bcb .eq. 1) then
                    beta=1
                else
                    beta=0
                end if

                ! materials
                m_L = mt( cix(i-1), ciy(j  ) )
                m_R = mt( cix(i+1), ciy(j  ) )
                m_T = mt( cix(i  ), ciy(j+1) )
                m_C = mt( cix(i  ), ciy(j  ) )

                A_L = -2.0*dc(m_L,g)*dc(m_C,g)*dy(j)/(dx(i-1)*dc(m_C,g)+dx(i)*dc(m_L,g))
                A_R = -2.0*dc(m_R,g)*dc(m_C,g)*dy(j)/(dx(i+1)*dc(m_C,g)+dx(i)*dc(m_R,g))
                A_T = -2.0*dc(m_T,g)*dc(m_C,g)*dx(i)/(dy(j+1)*dc(m_C,g)+dy(j)*dc(m_T,g))
                A_C = dv(k)*sr(m_C,g) - (A_L+A_R+A_T) &
                      + dc(m_C,g)*dx(i)*2*(1-beta)/(4*(1+beta)*dc(m_C,g)+dy(j)*(1-beta))

                call SetValues( A, inp, rhs, dat, k, lenk, g, (/A_L,A_R,A_B,A_T,A_C/), &
                                (/1,1,0,1/), dv(k), cix(i), ciy(j), m_C )

                !----------------------------------------------
                ! TOP EDGE
                !----------------------------------------------

                j = inp%nyfm
                k = i+(j-1)*( inp%nxfm )

                if (inp%bct .eq. 1) then
                    beta=1
                else
                    beta=0
                end if

                ! materials
                m_L = mt( cix(i-1), ciy(j  ) )
                m_R = mt( cix(i+1), ciy(j  ) )
                m_B = mt( cix(i  ), ciy(j-1) )
                m_C = mt( cix(i  ), ciy(j  ) )

                A_L = -2.0*dc(m_L,g)*dc(m_C,g)*dy(j)/(dx(i-1)*dc(m_C,g)+dx(i)*dc(m_L,g))
                A_R = -2.0*dc(m_R,g)*dc(m_C,g)*dy(j)/(dx(i+1)*dc(m_C,g)+dx(i)*dc(m_R,g))
                A_B = -2.0*dc(m_B,g)*dc(m_C,g)*dx(i)/(dy(j-1)*dc(m_C,g)+dy(j)*dc(m_B,g))
                A_C = dv(k)*sr(m_C,g) - (A_L+A_R+A_B) &
                      + dc(m_C,g)*dx(i)*2*(1-beta)/(4*(1+beta)*dc(m_C,g)+dy(j)*(1-beta))

                call SetValues( A, inp, rhs, dat, k, lenk, g, (/A_L,A_R,A_B,A_T,A_C/), &
                                (/1,1,1,0/), dv(k), cix(i), ciy(j), m_C )

            end do

            ! finally, the corners

            !----------------------------------------------
            ! BOTTOM LEFT
            !----------------------------------------------

            i=1
            j=1
            k = i+(j-1)*( inp%nxfm )
            if (inp%bcl .eq. 1) then
                beta1=1
            else
                beta1=0
            end if
            if (inp%bcb .eq. 1) then
                beta2=1
            else
                beta2=0
            end if

            ! materials
            m_R = mt( cix(i+1), ciy(j  ) )
            m_T = mt( cix(i  ), ciy(j+1) )
            m_C = mt( cix(i  ), ciy(j  ) )

            A_R = -2.0*dc(m_R,g)*dc(m_C,g)*dy(j)/(dx(i+1)*dc(m_C,g)+dx(i)*dc(m_R,g))
            A_T = -2.0*dc(m_T,g)*dc(m_C,g)*dx(i)/(dy(j+1)*dc(m_C,g)+dy(j)*dc(m_T,g))
            A_C = dv(k)*sr(m_C,g) - (A_R+A_T) &
                  + dc(m_C,g)*dy(j)*2*(1-beta1)/(4*(1+beta1)*dc(m_C,g)+dx(i)*(1-beta1)) &
                  + dc(m_C,g)*dx(i)*2*(1-beta2)/(4*(1+beta2)*dc(m_C,g)+dy(j)*(1-beta2))

            call SetValues( A, inp, rhs, dat, k, lenk, g, (/A_L,A_R,A_B,A_T,A_C/), &
                            (/0,1,0,1/), dv(k), cix(i), ciy(j), m_C )

            !----------------------------------------------
            ! BOTTOME RIGHT
            !----------------------------------------------

            i=1
            j=inp%nyfm
            k = i+(j-1)*( inp%nxfm )
            if (inp%bcl .eq. 1) then
                beta1=1
            else
                beta1=0
            end if
            if (inp%bct .eq. 1) then
                beta2=1
            else
                beta2=0
            end if

            ! materials
            m_R = mt( cix(i+1), ciy(j  ) )
            m_B = mt( cix(i  ), ciy(j-1) )
            m_C = mt( cix(i  ), ciy(j  ) )

            A_R = -2.0*dc(m_R,g)*dc(m_C,g)*dy(j)/(dx(i+1)*dc(m_C,g)+dx(i)*dc(m_R,g))
            A_B = -2.0*dc(m_B,g)*dc(m_C,g)*dx(i)/(dy(j-1)*dc(m_C,g)+dy(j)*dc(m_B,g))
            A_C = dv(k)*sr(m_C,g) - (A_R+A_B) &
                             + dc(m_C,g)*dy(j)*2*(1-beta1)/(4*(1+beta1)*dc(m_C,g)+dx(i)*(1-beta1)) &
                             + dc(m_C,g)*dx(i)*2*(1-beta2)/(4*(1+beta2)*dc(m_C,g)+dy(j)*(1-beta2))

            call SetValues( A, inp, rhs, dat, k, lenk, g, (/A_L,A_R,A_B,A_T,A_C/), &
                            (/0,1,1,0/), dv(k), cix(i), ciy(j), m_C )

            !----------------------------------------------
            ! TOP RIGHT
            !----------------------------------------------

            i=inp%nxfm
            j=1
            k = i+(j-1)*( inp%nxfm )
            if (inp%bcr.eq.1) then
                beta1=1
            else
                beta1=0
            end if
            if (inp%bcb.eq.1) then
                beta2=1
            else
                beta2=0
            end if

            ! materials
            m_L = mt( cix(i-1), ciy(j  ) )
            m_T = mt( cix(i  ), ciy(j+1) )
            m_C = mt( cix(i  ), ciy(j  ) )

            A_L = -2.0*dc(m_L,g)*dc(m_C,g)*dy(j)/(dx(i-1)*dc(m_C,g)+dx(i)*dc(m_L,g))
            A_T = -2.0*dc(m_T,g)*dc(m_C,g)*dx(i)/(dy(j+1)*dc(m_C,g)+dy(j)*dc(m_T,g))
            A_C = dv(k)*sr(m_C,g) - (A_L+A_T) &
                  + dc(m_C,g)*dy(j)*2*(1-beta1)/(4*(1+beta1)*dc(m_C,g)+dx(i)*(1-beta1)) &
                  + dc(m_C,g)*dx(i)*2*(1-beta2)/(4*(1+beta2)*dc(m_C,g)+dy(j)*(1-beta2))

            call SetValues( A, inp, rhs, dat, k, lenk, g, (/A_L,A_R,A_B,A_T,A_C/), &
                            (/1,0,0,1/), dv(k), cix(i), ciy(j), m_C )

            ! TOP RIGHT -----------------------------------------------------------
            i=inp%nxfm
            j=inp%nyfm
            k = i+(j-1)*( inp%nxfm )
            if (inp%bcr .eq. 1) then
                beta1=1
            else
                beta1=0
            end if
            if (inp%bct .eq. 1) then
                beta2=1
            else
                beta2=0
            end if

            ! materials
            m_L = mt( cix(i-1), ciy(j  ) )
            m_B = mt( cix(i  ), ciy(j-1) )
            m_C = mt( cix(i  ), ciy(j  ) )

            A_L = -2.0*dc(m_L,g)*dc(m_C,g)*dy(j)/(dx(i-1)*dc(m_C,g)+dx(i)*dc(m_L,g))
            A_B = -2.0*dc(m_B,g)*dc(m_C,g)*dx(i)/(dy(j-1)*dc(m_C,g)+dy(j)*dc(m_B,g))
            A_C = dv(k)*sr(m_C,g) - (A_L+A_B) &
                  + dc(m_C,g)*dy(j)*2*(1-beta1)/(4*(1+beta1)*dc(m_C,g)+dx(i)*(1-beta1)) &
                  + dc(m_C,g)*dx(i)*2*(1-beta2)/(4*(1+beta2)*dc(m_C,g)+dy(j)*(1-beta2))

            call SetValues( A, inp, rhs, dat, k, lenk, g, (/A_L,A_R,A_B,A_T,A_C/), &
                            (/1,0,1,0/), dv(k), cix(i), ciy(j), m_C )

            ! finalize the matrix assembly for this group
            call MatAssemblyBegin(A(g),MAT_FINAL_ASSEMBLY,ierr)
            call MatAssemblyEnd(A(g),MAT_FINAL_ASSEMBLY,ierr)

        end do ! group loop

    end subroutine setk

    !> set the matrix and vector values
    subroutine SetValues( A, inp, rhs, dat, k, lenk, g, val, ind, dv, cix_i, ciy_j, m_C )

        implicit none
#include "finclude/petsc.h90"
        Mat, dimension(:), pointer, intent(inout) :: A            ! groupwise matrix operators
        type(indat),   intent(inout)              :: inp          ! input data
        type(rhsdat),  intent(inout)              :: rhs          ! right hand side data
        type(coefdat), intent(inout)              :: dat          ! other data (absorption etc.)
        PetscInt,      intent(in)                 :: k            ! current row index (1-indexing)
        PetscInt,      intent(in)                 :: lenk         ! matrix size
        PetscInt,      intent(in)                 :: g            ! current group index (1-indexing)
        PetscScalar,   intent(in)                 :: val(:)       ! values to be entered
        PetscInt,      intent(in)                 :: ind(:)       ! e.g. (1 1 1 0) means include L,R,B but not T
        PetscScalar,   intent(in)                 :: dv           ! current volume
        PetscInt,      intent(in)                 :: cix_i, ciy_j ! source indices
        PetscInt,      intent(in)                 :: m_C          ! material placement index
        PetscInt                                  :: rownum       ! number of rows added
        PetscInt                                  :: rowidx(1)    ! row index (k-1 for 0-indexing)
        PetscInt                                  :: colnum       ! number of columns added
        PetscInt                                  :: colidx(5)    ! col index (0-indexed)
        PetscInt                                  :: ierr, &      ! error flag
                                                     gg           ! secondary group index
        ! define row index
        rownum    = 1
        rowidx(1) = k-1

        ! define column indices
        colnum    = 5
        colidx(1) = rowidx(1)-1 ! L
        colidx(2) = rowidx(1)+1 ! R
        colidx(3) = rowidx(1)-inp%nxfm ! B
        colidx(4) = rowidx(1)+inp%nxfm ! T
        colidx(5) = rowidx(1) ! C
        if (ind(1).eq.0) colidx(1) = -1
        if (ind(2).eq.0) colidx(2) = -1 ! setting to negative causes
        if (ind(3).eq.0) colidx(3) = -1
        if (ind(4).eq.0) colidx(4) = -1 ! Petsc to ignore it (perfect for edges)

        call MatSetValues( A(g), rownum, rowidx, colnum, colidx, val, INSERT_VALUES, ierr )
        do gg = 1, inp%numg
            call VecSetValues( rhs%sigScatter(g,gg), rownum, k-1, dv*inp%sc(m_C,g,gg), INSERT_VALUES, ierr )
        end do
        ! compute fixed source if needed
        if (inp%ptype .le. 0) then
            call VecSetValues( rhs%fixedSource(g), rownum, k-1, dv*inp%src(cix_i,ciy_j,g), INSERT_VALUES, ierr )
        end if
        call VecSetValues( rhs%nuSigF(g),  rownum, k-1, dv*inp%ns(m_C,g), INSERT_VALUES, ierr )
        call VecSetValues( rhs%xiSpect(g), rownum, k-1, inp%xi(m_C,g),    INSERT_VALUES, ierr )
        call VecSetValues( dat%sigAbs(g),  rownum, k-1, dv*inp%ab(m_C,g), INSERT_VALUES, ierr )

    end subroutine SetValues

end module matx_const
